import json
from langchain_core.messages import AIMessage
from typing import Dict, List
from langchain_core.runnables import Runnable, RunnablePassthrough
from langchain_core.tools import tool
from ChatGLM_new import tongyi_llm

@tool
def count_emails(last_n_days: int) -> int:
    """Multiply two integers together."""
    return last_n_days * 2


@tool
def send_email(message: str, recipient: str) -> str:
    "Add two integers."
    return f"Successfully sent email to {recipient}."


tools = [count_emails, send_email]
llm_with_tools = tongyi_llm.bind_tools(tools)


def call_tools(msg: AIMessage) -> List[Dict]:
    """Simple sequential tool calling helper."""
    tool_map = {tool.name: tool for tool in tools}
    tool_calls = msg.tool_calls.copy()
    for tool_call in tool_calls:
        tool_call["output"] = tool_map[tool_call["name"]].invoke(tool_call["args"])
    return tool_calls


class NotApproved(Exception):
    """自定义异常."""


def human_approval(msg: AIMessage) -> AIMessage:
    """负责传递其输入或引发异常.

    Args:
        msg: 聊天模型的输出

    Returns:
        msg: msg的原始输出
    """
    tool_strs = "\n\n".join(
        json.dumps(tool_call, indent=2) for tool_call in msg.tool_calls
    )
    input_msg = (
        f"您是否赞成以下工具调用\n\n{tool_strs}\n\n"
        "除此以外的任何内容 'Y'/'Yes'（不区分大小写）将被视为否。\n >>>"
    )
    resp = input(input_msg)
    if resp.lower() not in ("yes", "y"):
        raise NotApproved(f"工具调用未获批准:\n\n{tool_strs}")
    return msg

chain = llm_with_tools | human_approval | call_tools
res = chain.invoke("我在过去 5 天内收到了多少封电子邮件?")
print(res)